from ppgnss import gnss_utils
from datetime import datetime

import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.axes_grid1 import make_axes_locatable

import utils
import ttide
from ttide.t_tide import t_tide


plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']

width, height = gnss_utils.cm2inch(18), gnss_utils.cm2inch(8)

filename = "span5_comp_va_5min_down.txt"
data = utils.read_data_file(filename)
a_time_from, a_time_to = "2023-01-29 00:00:00", "2023-03-09 00:00:00" # 50 days

xr_neu = gnss_utils.loadobject("enu_0129_0309.npo")
data_seal = gnss_utils.loadobject("sea_level_0129_0309.npo")

xr_seal_down, xr_seal_up, xr_delta = data_seal["down"], data_seal["up"], data_seal["delta"]

xr_seal_down = xr_seal_down.loc[a_time_from:a_time_to]
xr_seal_up = xr_seal_up.loc[a_time_from:a_time_to]
xr_delta = xr_delta.loc[a_time_from:a_time_to]
width, height = gnss_utils.cm2inch(18), gnss_utils.cm2inch(12)

dirs = ["u", "v", "a"]
xyznames = ["U", "Y", "X"]
fig, axes = plt.subplots(nrows=5, ncols=1, sharey=False, sharex="col", figsize=(width, height))
plt.subplots_adjust(left=0.1, right=0.98, bottom=0.1, top=0.98, hspace=0)

y_lim_min, y_lim_max = -0.02, 0.02

for icom, com in enumerate(dirs):
    obs1 = xr_neu.loc[:, com]
    stime = datetime.strptime(str(obs1.coords["time"].values[0])[:19], '%Y-%m-%dT%H:%M:%S')
    print(stime)
    obs1 = obs1 - obs1.mean()
    # obs1 = obs1*10
    # obs1 = obs1.interpolate_na(dim="time")
    out1 = t_tide(obs1.values, dt=5/60., stime=stime, lat=np.array(30.23), synth=1)
    # axes[icom].plot(obs1.coords["time"], obs1, linestyle='-', alpha=0.7, label=com.upper())
    # axes[icom].set_ylim((y_lim_min, y_lim_max))
    # axes[icom].plot(obs1.coords["time"], out1["xout"].squeeze(), linestyle='--', linewidth=1, label=u'Prediction')
    # axes[icom][0].plot(obs1.coords["time"], obs1-out1["xout"].squeeze(), label=u'Error')
    # axes[icom][0].fill_between(obs1.coords["time"], obs1-out1["xout"].squeeze(), where=(obs1-out1["xout"].squeeze()< 0), color="gray", alpha=0.5)
    # axes[icom][0].fill_between(obs1.coords["time"], obs1-out1["xout"].squeeze(), where=(obs1-out1["xout"].squeeze()> 0), color="gray", alpha=0.5)
    # axes[icom].legend(numpoints=1, loc='lower right')
    # axes[icom].grid(axis='y', linestyle='--')
    # axes[icom].set_yticks(np.arange(-0.02, 0.03, 0.005))
    fs = 60/5  # 采样频率，1小时采样一次
    signal = obs1.values
    signal = np.nan_to_num(signal, copy=True, nan=0)
    fft_result = np.fft.fft(signal)
    fft_freqs = np.fft.fftfreq(len(fft_result), 1/fs)  # 计算频率轴
# 提取正频率分量
    positive_freq_mask = fft_freqs > 0
    fft_freqs_positive = fft_freqs[positive_freq_mask]
    fft_result_positive = fft_result[positive_freq_mask] 
    axes[icom].plot(fft_freqs_positive, 2*np.abs(fft_result_positive)/(len(signal)/2), c="k", alpha=1, linewidth=1)
    axes[icom].set_ylim((0, 0.003))
    axes[icom].set_ylabel(xyznames[icom].upper())
    axes[icom].set_yticks([0, 0.001, 0.002], ["0", "1 mm", "2 mm"], fontsize=10)
    for freq, name in zip(data["Ufreq"], data["Utide"]):
        axes[icom].axvline(x=freq, color='grey', linestyle='--', linewidth=0.4, label=name)
        # axes[icom][1].text(freq, 0.2, name, fontsize=10)


# obs = xr_delta - xr_delta.mean()
obs = xr_seal_down - xr_seal_down.mean()
stime = datetime.strptime(str(obs.coords["time"].values[0])[:19], '%Y-%m-%dT%H:%M:%S')
out = t_tide(obs, dt=5/60, stime=stime, lat=np.array(30.23), synth=1)
# axes[3].plot(obs.coords["time"], obs, label=u'Sea Level')
# obs2 = xr_seal_up - xr_seal_up.mean()
# axes[3].plot(obs.coords["time"], out["xout"].squeeze(), label=u'Prediction')
# axes[3].set_ylim((-5, 5))
# axes[3].plot(xr_seal_up.coords["time"], xr_seal_up, label=u'Up Sea Level', c="r")
# axes[3].plot(xr_seal_down.coords["time"], xr_seal_down, label=u'Down Sea Level', c="k")
# axes[3].legend(numpoints=1, loc='lower right')
fs = 60/5  # 采样频率，1小时采样一次
signal = obs.values
signal = np.nan_to_num(signal, nan=0)
signal = np.hamming(len(signal))*signal
fft_result = np.fft.fft(signal)
fft_freqs = np.fft.fftfreq(len(fft_result), 1/fs)  # 计算频率轴
# 提取正频率分量
positive_freq_mask = fft_freqs > 0
fft_freqs_positive = fft_freqs[positive_freq_mask]
fft_result_positive = fft_result[positive_freq_mask]
axes[3].plot(fft_freqs_positive, 2*np.abs(fft_result)[positive_freq_mask]/(len(signal)/2), c="k", alpha=1, linewidth=1)
axes[3].set_ylim((0, 2.5))
axes[3].set_ylabel("闸外水位")
axes[3].set_yticks([0, 1, 2], ["0", "1.0 m", "2.0 m"], fontsize=10)
for freq, name in zip(data["Ufreq"], data["Utide"]):
    axes[3].axvline(x=freq, color='grey', linestyle='--', linewidth=0.4,  label=name)
    # axes[3][1].text(freq, 100, name, fontsize=10)
# axes[3].set_title('Frequency Spectrum')
# axes[3].set_xlabel('Frequency (cycles per hour)')

obs = xr_delta - xr_delta.mean()
# obs = xr_seal_down - xr_seal_down.mean()
stime = datetime.strptime(str(obs.coords["time"].values[0])[:19], '%Y-%m-%dT%H:%M:%S')
out = t_tide(obs, dt=5/60, stime=stime, lat=np.array(30.23), synth=1)
# axes[4].plot(obs.coords["time"], obs, label=u'Sea Level')
# obs2 = xr_seal_up - xr_seal_up.mean()
# axes[4].plot(obs.coords["time"], out["xout"].squeeze(), label=u'Prediction')
# axes[4].set_ylim((-5, 5))
# axes[3].plot(xr_seal_up.coords["time"], xr_seal_up, label=u'Up Sea Level', c="r")
# axes[3].plot(xr_seal_down.coords["time"], xr_seal_down, label=u'Down Sea Level', c="k")
# axes[4].legend(numpoints=1, loc='lower right')
fs = 60/5 # 采样频率
signal = obs.values
signal = np.nan_to_num(signal, nan=0)
signal = np.hamming(len(signal))*signal
fft_result = np.fft.fft(signal)
fft_freqs = np.fft.fftfreq(len(fft_result), 1/fs)  # 计算频率轴
# 提取正频率分量
positive_freq_mask = fft_freqs > 0
fft_freqs_positive = fft_freqs[positive_freq_mask]
fft_result_positive = fft_result[positive_freq_mask]
axes[4].plot(fft_freqs_positive, 2*np.abs(fft_result)[positive_freq_mask]/(len(signal)/2), c="k", alpha=1, linewidth=1)
axes[4].set_ylim((0, 2.5))
axes[4].set_ylabel("内外位差")
axes[4].set_yticks([0, 1, 2], ["0", "1.0 m", "2.0 m"], fontsize=10)
for freq, name in zip(data["Ufreq"], data["Utide"]):
    axes[4].axvline(x=freq, color='grey', linestyle='--', linewidth=0.4,  label=name)
    # axes[3][1].text(freq, 100, name, fontsize=10)
# axes[4].set_title('Frequency Spectrum')
axes[4].set_xlabel(u'频率 (周/小时)', fontsize=10)
plt.xlim((0, 0.25))
fig_filename = "figures/fftfreqs.png"
plt.savefig(fig_filename, dpi=300)
plt.show() 

